from typing_extensions import Self, Any
from utils.db import async_session_maker
from utils.repositories import (UserRepository,
                                StateRepository,
                                RegionRepository,
                                FlightRepository)


class UnitOfWork:
    session_factory = async_session_maker

    async def __aenter__(self) -> Self:
        self.session = self.session_factory()

        self.users = UserRepository(self.session)
        self.states = StateRepository(self.session)
        self.regions = RegionRepository(self.session)
        self.flights = FlightRepository(self.session)

        return self

    async def __aexit__(self, exc_type: Exception, *args: Any) -> None:
        if exc_type is None:
            await self.commit()
        else:
            await self.rollback()

        await self.session.close()
        return

    async def commit(self) -> None:
        await self.session.commit()
        return

    async def rollback(self) -> None:
        await self.session.rollback()
        return
